package com.atomgit.ai.proxy.processor;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.atomgit.ai.proxy.param.PictureQuesion;
import com.atomgit.ai.proxy.param.PromptParam;
import com.atomgit.ai.proxy.param.inner.tongyivl.TongYiVLParam;
import com.atomgit.ai.proxy.param.inner.tongyivl.TongYiVLParamInput;
import com.atomgit.ai.proxy.param.inner.tongyivl.TongYiVLParamInputMessages;
import com.atomgit.ai.proxy.processor.service.ParamAndResponseProcessor;
import com.atomgit.ai.proxy.vo.ModelAnswerOutPutVO;
import com.atomgit.ai.proxy.vo.ModelAnswerVO;
import lombok.extern.slf4j.Slf4j;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 通义千问VL参数与结果处理类
 */
@Slf4j
public class ParamAndResponseProcessorForTongYiVL implements ParamAndResponseProcessor{
    @Override
    public String generateRequestBody(PromptParam param) {
        TongYiVLParam tongYiVLParam = new TongYiVLParam();
        tongYiVLParam.setModel(param.getVersion());
        TongYiVLParamInput tongYiVLParamInput = new TongYiVLParamInput();
        List<PictureQuesion> pictureQuesionList = param.getPictureQuesionList();
        List<TongYiVLParamInputMessages> messages = new ArrayList<>();

        for (int i = 0; i < pictureQuesionList.size(); i++) {
            List<HashMap<String, String>>  content = new ArrayList<>();
            PictureQuesion pictureQuesion =  pictureQuesionList.get(i);
            TongYiVLParamInputMessages tongYiVLParamInputMessages = new TongYiVLParamInputMessages();
            tongYiVLParamInputMessages.setRole(pictureQuesion.getRole());
            if (pictureQuesion.getImage() != null && !pictureQuesion.getImage().trim().equals("")){
                HashMap<String, String> imageContent = new HashMap<>();
                imageContent.put("image", pictureQuesion.getImage());
                content.add(imageContent);
            }
            if (pictureQuesion.getText() != null && !pictureQuesion.getText().trim().equals("")){
                HashMap<String, String> textContent = new HashMap<>();
                textContent.put("text", pictureQuesion.getText());
                content.add(textContent);
            }
            tongYiVLParamInputMessages.setContent(content);
            messages.add(tongYiVLParamInputMessages);
        }
        tongYiVLParamInput.setMessages(messages);
        tongYiVLParam.setInput(tongYiVLParamInput);
        return JSON.toJSONString(tongYiVLParam);
    }

    @Override
    public ModelAnswerVO generateModelAnswerVO(String responseBody) {
        if (responseBody == null) {
            return null;
        }
        log.info("responseBody:"+responseBody);
        ModelAnswerVO modelAnswerVO = new ModelAnswerVO();
        ModelAnswerOutPutVO output = new ModelAnswerOutPutVO();
        output.setFinish_reason("stop");
        JSONObject jsonObject = JSON.parseObject(responseBody);
        String text = jsonObject.getJSONObject("output").getJSONArray("choices").getJSONObject(0).getJSONObject("message").getString("content");
        output.setText(text);
        modelAnswerVO.setOutput(output);
        return modelAnswerVO;
    }

}
